package com.example.springboot.utils;

import javax.websocket.Session;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * @Description: VoteSocketUtils
 * @Author wenlinyu
 * @Date: 2020/12/21 11:10
 * @Version 1.0
 */
public class VoteSocketUtils {

    private static Map<String, CopyOnWriteArraySet<Session>> voteSocketMap = new ConcurrentHashMap<>();

    public static void put(String voteId, Session session) {
        CopyOnWriteArraySet<Session> voteSocketSet = getVoteSocketSet(voteId);
        voteSocketSet.add(session);
        voteSocketMap.put(voteId, voteSocketSet);
    }

    public static void remove(String voteId, Session session) {
        CopyOnWriteArraySet<Session> voteSocketSet = getVoteSocketSet(voteId);
        voteSocketSet.remove(session);
    }

    public static void sendMessageAll(String voteId, String message) throws IOException {
        CopyOnWriteArraySet<Session> voteSocketSet = getVoteSocketSet(voteId);
        for (Session session : voteSocketSet) {
            session.getBasicRemote().sendText(message);
        }
    }

    public static CopyOnWriteArraySet<Session> getVoteSocketSet(String voteId) {
        CopyOnWriteArraySet<Session> voteSocketSet = voteSocketMap.get(voteId);
        if (voteSocketSet == null) {
            voteSocketSet = new CopyOnWriteArraySet<>();
        }
        return voteSocketSet;
    }

}
